"""
使用flask部署pytorch模型
"""

import io

from flask import Flask
import flask
import torchvision
from PIL import Image
import torch

from model import Net

app = Flask(__name__)
model = None
use_gpu = True
pre_transforms = torchvision.transforms.Compose([
    # torchvision.transforms.ToPILImage(),
    # torchvision.transforms.SSDCropping(),
    # torchvision.transforms.Resize(size=(256, 256)),
    torchvision.transforms.Resize(size=(28, 28)),
    # torchvision.transforms.ColorJitter(),
    torchvision.transforms.ToTensor(),
    # torchvision.transforms.RandomHorizontalFlip(),
    # torchvision.transforms.Normalization(),
    torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),  # 归一化
    # torchvision.transforms.AssignGTtoDefaultBox()
])
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


@app.route('/')
def hello_world():
    return 'Hello, World!'


def load_model():
    """
    Load the pre-trained model, you can use your model just as easily.
    """
    global model
    model = Net()
    model.load_state_dict(torch.load('model.pth', map_location=device))
    model.eval()
    if use_gpu:
        model = model.to(device)


load_model()


def prepare_image(img0):
    img0 = pre_transforms(img0)
    return img0


@app.route("/predict", methods=["POST"])
def predict():
    data = {"success": False}
    global model
    if flask.request.method == 'POST':
        if flask.request.files.get("image"):
            img0 = flask.request.files["image"].read()
            img0 = Image.open(io.BytesIO(img0))
            img0 = prepare_image(img0=img0)
            img0 = img0.unsqueeze(0).to(device)

            # img0 = Image.fromarray(np.uint8(img0))

            # posture = np.array(network(img0))  # .detach().numpy()
            posture = model(img0)  # .detach().numpy()
            print('posture:', posture)
            data['predictions'] = str(posture)
            data["success"] = True

    return flask.jsonify(data)


if __name__ == '__main__':
    load_model()
    app.run(debug=True, host='0.0.0.0', port=5000)
